#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from geometry_msgs.msg import Twist, Vector3
import numpy as np
import sys
import termios
import tty
import threading

class BalanceRobot(Node):
    def __init__(self):
        super().__init__('balance_robot')
        
        # 移除测试模式
        self.test_mode = False
        
        # 设置目标角度为-78度（转换为弧度）
        self.target_angle = np.radians(-78.0)
        
        # 添加全局速度比例（5%）
        self.speed_ratio = 0.05
        
        # 初始化当前角度
        self.current_angle = self.target_angle
        self.imu_ready = False
        
        # 创建euler_angles订阅者（替代IMU）
        self.euler_sub = self.create_subscription(
            Vector3,
            '/euler_angles',
            self.euler_callback,
            10)
            
        # 创建发布者
        self.cmd_vel_pub = self.create_publisher(
            Twist,
            '/cmd_vel',
            10)
            
        # PID控制器参数
        self.Kp = 10.0
        self.Ki = 0.0
        self.Kd = 2.0
        
        # PID控制器状态
        self.prev_error = 0.0
        self.integral = 0.0
        
        # 运动控制变量
        self.linear_speed = 0.0
        self.angular_speed = 0.0
        
        # 创建定时器进行控制
        self.timer = self.create_timer(0.01, self.control_loop)  # 100Hz控制频率

    def euler_callback(self, msg):
        # euler_angles中的y值对应pitch角度
        self.current_angle = msg.y
        self.imu_ready = True
        # 显示当前Pitch角度（度数）
        pitch_deg = np.degrees(self.current_angle)
        # 使用print而不是ROS logger，确保输出更清晰
        print(f'\r当前Pitch角度: {pitch_deg:.2f}°, 目标角度: {np.degrees(self.target_angle):.2f}°', end='', flush=True)

    def pid_control(self):
        error = self.target_angle - self.current_angle
        self.integral += error * 0.01  # dt = 0.01s
        derivative = (error - self.prev_error) / 0.01
        
        output = self.Kp * error + self.Ki * self.integral + self.Kd * derivative
        self.prev_error = error
        
        return output

    def control_loop(self):
        try:
            # 等待IMU数据准备就绪
            if not self.imu_ready:
                return
                
            # 计算平衡控制输出
            balance_output = self.pid_control()
            
            # 创建并发布控制命令，应用全局速度比例
            cmd_vel = Twist()
            cmd_vel.linear.x = (self.linear_speed + balance_output) * self.speed_ratio
            cmd_vel.angular.z = self.angular_speed * self.speed_ratio
            self.cmd_vel_pub.publish(cmd_vel)
            
        except Exception as e:
            self.get_logger().error(f'Control loop error: {str(e)}')

    def stop_robot(self):
        """安全停止机器人"""
        stop_cmd = Twist()
        stop_cmd.linear.x = 0.0
        stop_cmd.angular.z = 0.0
        # 多发送几次确保停止命令被执行
        for _ in range(5):
            self.cmd_vel_pub.publish(stop_cmd)

    def update_motion_command(self, linear, angular):
        self.linear_speed = linear
        self.angular_speed = angular

def get_key():
    settings = termios.tcgetattr(sys.stdin)
    tty.setraw(sys.stdin.fileno())
    key = sys.stdin.read(1)
    termios.tcsetattr(sys.stdin, termios.TCSADRAIN, settings)
    return key

def key_control(node):
    while True:
        key = get_key()
        
        if key == 'w':
            node.linear_speed = 0.2
        elif key == 's':
            node.linear_speed = -0.2
        elif key == 'a':
            node.angular_speed = 0.5
        elif key == 'd':
            node.angular_speed = -0.5
        elif key == 'q':
            break
        else:  # 其他键停止运动
            node.linear_speed = 0.0
            node.angular_speed = 0.0

def main(args=None):
    rclpy.init(args=args)
    balance_robot = BalanceRobot()
    
    # 创建键盘控制线程
    key_thread = threading.Thread(target=key_control, args=(balance_robot,))
    key_thread.daemon = True  # 设置为守护线程，这样主线程结束时它会自动结束
    key_thread.start()
    
    try:
        rclpy.spin(balance_robot)
    except KeyboardInterrupt:
        print("\n正在安全停止机器人...")
        balance_robot.stop_robot()
    finally:
        balance_robot.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()
